
import sys
import os
import os.path
import re
import struct
from cStringIO import StringIO
from array import array
import zlib

def need_escape(s):
	encoding = sys.getfilesystemencoding()
	try:
		return s != s.encode(encoding).decode(encoding)
	except:
		return True

def encode_char(m):
	if isinstance(m, basestring): # FIXME: type check is not good
		c = m
	else:
		c = m.group()
	if len(c) == 1:
		return '_u%04x' % ord(c)
	else:
		assert re.match(r'_u[0-9a-fA-F]{4}', c)
		return encode_char(c[0])+c[1:]

def force_encode_path(path):
	return re.sub(r'_u[0-9a-fA-F]{4}|[^-\w\s\\/.]', encode_char, path)

def encode_path(path):
	if need_escape(path):
		return force_encode_path(path)
	else:
		return path

def decode_char(m):
	chars = m.group()
	#assert re.match(r'_u[0-9a-fA-F]{4}', chars)
	return unichr(int(chars[2:], 16))

def decode_path(path):
	return re.sub(r'_u[0-9a-fA-F]{4}', decode_char, path)

def file_size(path):
	return os.path.getsize(path)
	return os.stat(path).st_size

def skip_string(x, string):
	assert x.read(len(string)) == string

def read_value(x, format):
	return struct.unpack(format, x.read(struct.calcsize(format)))[0]

def read_values(x, format):
	return struct.unpack(format, x.read(struct.calcsize(format)))

def read_uint(x):
	#return read_value(x, 'I')
	return struct.unpack('I', x.read(4))[0]

def read_ushort(x):
	return read_value(x, 'H')

def read_ulong(x):
	#return read_value(x, 'Q')
	return struct.unpack('Q', x.read(8))[0]

def read_string(x):
	length = read_ushort(x)
	return x.read(length*2).decode('UTF-16LE')

class Segment:
	def __init__(self, stream):
		self.compressed = read_uint(stream)
		self.offset = read_ulong(stream)
		self.original_size = read_ulong(stream)
		self.compressed_size = read_ulong(stream)

class Entry:
	def __init__(self, stream):
		'''
		'''
		skip_string(stream, 'File')
		entry_size = read_ulong(stream)

		skip_string(stream, 'info')
		info_size = read_ulong(stream)

		self.encrypted = bool(read_uint(stream))
		self.original_size = read_ulong(stream)
		self.compressed_size = read_ulong(stream)

		self.name = read_string(stream)

		assert info_size == len(self.name) * 2 + 22
		skip_string(stream, 'segm')
		segment_number = read_ulong(stream) // 28   # 28 bytes per seg.

		self.segments = [Segment(stream) for i in range(segment_number)]

		assert sum(s.compressed_size for s in self.segments) == self.compressed_size
		assert sum(s.original_size for s in self.segments) == self.original_size

		skip_string(stream,'adlr')
		assert read_ulong(stream) == 4
		self.adler = read_uint(stream)
		assert entry_size == len(self.name) * 2 + segment_number * 28 + 62

	def encoded_name(self):
		return encode_path(self.name)

	def __repr__(self):
		return (u'<Entry %s>' % encode_path(self.name)).encode(sys.getfilesystemencoding())

def read_data_stream(x, entry, decrypter=None):
	buffer = StringIO()
	adler = zlib.adler32('')
	for segment in entry.segments :
		x.seek(segment.offset)
		if segment.compressed:
			data = zlib.decompress(x.read(segment.compressed_size))
		else:
			data = x.read(segment.compressed_size)
		assert len(data) == segment.original_size
		buffer.write(data)
		adler = zlib.adler32(data, adler)
	if ((adler + 0x0100000000L) & 0x00FFFFFFFFL) != entry.adler:
		# Convert to unsigned 32-bit integer
		print 'WARN: Checksum error (%s, %s), but continuing...' % (bin(adler), bin(entry.adler))
	data = buffer.getvalue()
	if entry.encrypted:
		assert decrypter
		return decrypter(data)
	else:
		if decrypter:
			return decrypter(data)
	return data

def read_data(path, entry, decrypter=None):
	with open(path, 'rb') as stream:
		return read_data_stream(stream, entry, decrypter)

def read_xp3(path):
	with open(path, 'rb') as x:
		skip_string(x, 'XP3\r\n \n\x1a\x8bg\x01')

		index_offset = read_ulong(x)
		assert index_offset < file_size(path)
		x.seek(index_offset)
		skip_string(x, '\x01')
		compressed_size, original_size = read_values(x, 'QQ')
		assert index_offset + compressed_size + 17 == file_size(path)
		uncompressed = x.read(compressed_size).decode('zlib')
		assert len(uncompressed) == original_size

	index_buffer = StringIO(uncompressed) # TODO: use with

	entries = []
	while index_buffer.tell() < original_size:
		entry = Entry(index_buffer)
		entries.append(entry)
		#print entry.name
		#print encode_path(entry.name)

	return entries

def fate_decrypter(raw):
	data = array('B', raw)
	for i in xrange(len(data)):
		data[i] ^= 0x0036
	if len(data) > 19:
		data[19] ^= 1
	if len(data) > 0x2ea29:
		data[0x2ea29] ^= 3
	return data.tostring()

try:
	import ctypes
	from common_c import load_library
	libxp3_fate = load_library('libxp3_fate')
	def fate_decrypter_c(raw):
		buffer = ctypes.create_string_buffer(raw, len(raw))
		libxp3_fate.xp3_fate_decrypter(buffer, len(raw))
		return buffer.raw
	fate_decrypter = fate_decrypter_c
except Exception:
	# failed to load, log an error?
	pass

def extract_xp3_to(xp3, folder, decrypter=None):
	entries = read_xp3(xp3)
	with open(xp3, 'rb') as x:
		if not os.path.exists(folder):
			os.makedirs(folder)
		for entry in entries:
			path = entry.encoded_name()
			if len(path) > 200:
				print 'WARN: ignoring long name item', path
				if entry.original_size != 0:
					print 'WARN: but file is not empty (size %d)' % entry.original_size
				continue
			dir = os.path.dirname(path)
			if dir:
				full_dir = os.path.join(folder, dir)
				if not os.path.exists(full_dir):
					os.makedirs(full_dir)
			full_path = os.path.join(folder, path)
			print full_path
			with open(full_path, 'wb') as output:
				output.write(read_data_stream(x, entry, decrypter))


def to_gbk(folder):
	def visitor(x, dirname, files):
		for f in files:
			p = os.path.join(dirname, f)
			if os.path.isfile(p) and re.match(r'.+\.(ks|fcf|tjs|txt)', p):
				with open(p, 'rb') as f:
					all = f.read()

				if re.search(r'^\xff\xfe', all):
					all = all.decode('utf-16le')[1:].encode('cp936')
				else:
					all = all.decode('cp932').encode('cp936', 'ignore')
				with open(p, 'wb') as f:
					f.write(all)

	os.path.walk(folder, visitor, None)


if __name__ == '__main__':
	pass
	#extract_xp3_to('data.xp3', 'data', fate_decrypter)
	#to_gbk('data')

	#for x in read_xp3('patch1.xp3'):
	#	print x.encoded_name()


